import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig
from functools import reduce


@GENERATOR_REGISTRY.register("ascend_generate_aclnn_cross_entropy_loss")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.input_dtype = None
        self.input_shape = []


    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]],
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        ## 所有输入tensor的dtype保持一致
        if index == 0:
            self.input_dtype = input_case.dtype # input dtype
            input_case.shape[0]=random.randint(1,8000)
            self.input_shape = input_case.shape # input shape [N, C]
            
        ## target与input的第0维保持一致
        if index == 1:
            input_case.shape[0] = self.input_shape[0] # target [N,]; 值域 [0, C) int64
            input_case.range_values = [-1, self.input_shape[1] - 1]

        ## weight与input的第1维保持一致
        if index == 2:
            input_case.shape[0] = self.input_shape[1]

        if index == 4:
            input_case.range_values = [-1, self.input_shape[1] - 1]

        return input_case